# 生成多种失真
import os
import glob
import numpy as np
from PIL import Image
import cv2
from tid2013_generator import tid2013_generator
from image_utils import open_image

def main():
    # 获取原始图像列表
    # 支持的扩展名（不区分大小写）
    exts = ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tif', '.tiff']
    # 获取所有文件
    all_files = glob.glob('./pristine_images/*')
    # 只保留扩展名匹配的文件
    file_list = [f for f in all_files if os.path.splitext(f)[1].lower() in exts]
    
    # 定义失真类型和级别
    distortions = [1, 2, 5, 6, 7, 8, 9, 10, 11, 14, 15, 16, 17, 18, 19, 22, 23]
    levels = [1, 2, 3, 4]
    
    for i, file_path in enumerate(file_list):
        # 读取原始图像
        ref_img = open_image(file_path)
        
        # 获取文件名信息
        filename = os.path.basename(file_path)
        
        # 为每种失真类型和级别生成失真图像
        for dist_type in distortions:  # 您可以选择要生成的失真类型
            for level in levels:  # 您可以决定生成多少级别的失真
                tid2013_generator(ref_img, dist_type, level, filename)
        
        print(f'完成图像 {i+1}*17 / {len(file_list)}*17...')

if __name__ == "__main__":
    main()